import torch

model = torch.load("model_0001900.pth")

del model['iteration']
del model['scheduler']
del model['optimizer']

del model["model"]["roi_heads.box.predictor.cls_score.weight"]
del model["model"]["roi_heads.box.predictor.cls_score.bias"]
# del model["model"]["module.roi_heads.box.predictor.bbox_pred.weight"]
# del model["model"]["module.roi_heads.box.predictor.bbox_pred.bias"]

torch.save(model, "modified_model.pth")
